import _thread as thread
import base64
import hashlib
import hmac
import json
import ssl
from datetime import datetime
from time import mktime
from urllib.parse import urlencode, urlparse
from wsgiref.handlers import format_date_time
import websocket


# =========================================
# 设置参数、调试prompt、定义接口
class Ws_Param(object):
    # 初始化
    def __init__(self):
        self.APIKey = "8c2334478d6a0c849dcc789972e3f785"
        self.APISecret = "OWQ5ZDJjOTA2NTdhY2NmNWU3NDkwOGIy"
        self.Spark_url = "wss://maas-api.cn-huabei-1.xf-yun.com/v1.1/chat"
        self.host = urlparse(self.Spark_url).netloc
        self.path = urlparse(self.Spark_url).path

    # 生成url
    def create_url(self):
        # 生成RFC1123格式的时间戳
        now = datetime.now()
        date = format_date_time(mktime(now.timetuple()))

        # 拼接字符串
        signature_origin = "host: " + self.host + "\n"
        signature_origin += "date: " + date + "\n"
        signature_origin += "GET " + self.path + " HTTP/1.1"

        # 进行hmac-sha256进行加密
        signature_sha = hmac.new(
            self.APISecret.encode("utf-8"),
            signature_origin.encode("utf-8"),
            digestmod=hashlib.sha256,
        ).digest()

        signature_sha_base64 = base64.b64encode(signature_sha).decode(
            encoding="utf-8")

        authorization_origin = f'api_key="{self.APIKey}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature_sha_base64}"'

        authorization = base64.b64encode(
            authorization_origin.encode("utf-8")).decode(
            encoding="utf-8"
        )

        # 将请求的鉴权参数组合为字典
        v = {"authorization": authorization, "date": date, "host": self.host}
        # 拼接鉴权参数，生成url
        url = self.Spark_url + "?" + urlencode(v)
        # 此处打印出建立连接时候的url,参考本demo的时候可取消上方打印的注释，比对相同参数时生成的url与自己代码生成的url是否一致
        return url


class DeepSeekMain(object):
    domain = "xdeepseekv3"
    appid = "275f9e51"
    answer = ""
    chat_prompt = \
    """
        [角色]：你叫灵犀，是鼎桥研发的四足机器人，搭载了国产机器人大脑软件硬件平台，可以和用户进行对话交流。
        
        [任务]：你需要准确回答用户提出的问题。
        
        [要求]：回答内容要用纯文本，并且回答长度不要超过50个字。标点符号只能用逗号和句号。 
        
        [边界]：禁止涉及暴力、隐私等问题。
    """
  
    cmd_prompt = \
    """
        # 角色
        
        你是一个智能机器人，能够理解人类的意图，并根据意图指令表示转成特定的指令。有指令冲突，不需要解释。
        
        #意图指令表
        
        | 序号 | 意图描述    | 输出指令                  |
        |------|-------------|---------------------------|
        | 1    | 起立        | up                        |
        | 2    | 趴下,休息,  | down                      |
        | 3    | 踏步        | step                      |
        | 4    | 前进        | forward                   |
        | 5    | 后退        | back                      |
        | 6    | 向左        | left                      |
        | 7    | 向右        | right                     |
        | 9    | 如果问到你看到了什么 | vision           |
        | 10   | 打招呼/自我介绍      | hello            |
        | 11   | 导航        | navigation-目的地<br>（例：navigation-实验室） |
        
    """

    def __int__(self):
        pass

    @staticmethod
    # 收到websocket错误的处理
    def on_error(ws, error):
        print("### error:", error)

    @staticmethod
    # 收到websocket关闭的处理
    def on_close(ws, one, two):
        print("\nws close ")

    @staticmethod
    # 收到websocket连接建立的处理
    def on_open(ws):
        print("\nws open ")
        thread.start_new_thread(DeepSeekMain.run, (ws,))

    @staticmethod
    def run(ws, *args):
        data = json.dumps(
            DeepSeekMain.gen_params(appid=ws.appid, domain=ws.domain,
                                    question=ws.question)
        )
        ws.send(data)

    # 收到websocket消息的处理
    @staticmethod
    def on_message(ws, message):
        data = json.loads(message)
        code = data["header"]["code"]
        if code != 0:
            print(f"请求错误: {code}, {data}")
            ws.close()
        else:
            choices = data["payload"]["choices"]
            status = choices["status"]
            content = choices["text"][0]["content"]
            DeepSeekMain.answer += content
            if status == 2:
                ws.close()

    @staticmethod
    def gen_params(appid, domain, question):
        """
        通过appid和用户的提问来生成请参数
        """
        data = {
            "header": {
                "app_id": appid,
                "uid": "1234",
                # "patch_id": ["XXXXXXXXXXXXXXXXXX"]
                # #调用微调大模型时必传, 否则不传。对应resourceId
            },
            "parameter": {
                "chat": {
                    "domain": domain,
                    "temperature": 0.5,
                    "max_tokens": 4096,  # 请根据不同模型支持范围，适当调整该值的大小
                }
            },
            "payload": {"message": {"text": question}},
        }
        return data

    @staticmethod
    def main(question, user_message=None):
        wsParam = Ws_Param()
        websocket.enableTrace(False)
        wsUrl = wsParam.create_url()
        ws = websocket.WebSocketApp(
            wsUrl,
            on_message=DeepSeekMain.on_message if not user_message else user_message,
            on_error=DeepSeekMain.on_error,
            on_close=DeepSeekMain.on_close,
            on_open=DeepSeekMain.on_open,
        )
        ws.question = question
        ws.appid = DeepSeekMain.appid
        ws.domain = DeepSeekMain.domain
        ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})


    def chat(text, user_message=None):
        question = [
            {
                "role": "system",
                "content": DeepSeekMain.chat_prompt,
            },
            {
                "role": "user",
                "content": text,
            },
        ]
        DeepSeekMain.answer = ""
        DeepSeekMain.main(question, user_message)
        return DeepSeekMain.answer.strip()


    def cmd(text, user_message=None):
        question = [
            {
                "role": "system",
                "content": DeepSeekMain.cmd_prompt,
            },
            {
                "role": "user",
                "content": text,
            },
        ]
        DeepSeekMain.answer = ""
        DeepSeekMain.main(question, user_message)
        return DeepSeekMain.answer.strip()
 